package IC.Context;

import IC.AST.*;
import IC.Symbols.Kind;
import IC.Symbols.Symbol;
import IC.Types.TypeTable;

import java.util.*;

/**
 * Created by IntelliJ IDEA.
 * User: nimrod
 * Date: 16/01/12
 * Time: 14:04
 * To change this template use File | Settings | File Templates.
 */

/*
    This visitor traverses the AST and makes sure that every variable usage has a preceding initialization.
    Using the $ret variable, it also checks if a method that has a return type, returns a value on all code paths.
 */
public class InitializationTestVisitor implements Visitor {

    private int currentInitDepth;
    private Object currentInitContext;
    private Set<Symbol> currentInitializedSymbols;

    private TypeTable typeTable;
    private List<ContextWarning> warnings;

    public InitializationTestVisitor(TypeTable typeTable) {
        this.typeTable = typeTable;
        warnings = new ArrayList<ContextWarning>();
    }

    private void setSymbolInit(Symbol sym) {
        if (this.currentInitDepth <= sym.getInitDepth()) {
            sym.setInitContext(this.currentInitContext);
            sym.setInitDepth(this.currentInitDepth);
            this.currentInitializedSymbols.add(sym);
        }
    }
	
	private void resetSymbolInit(Symbol sym) {
		sym.setInitContext(null);
		sym.setInitDepth(Integer.MAX_VALUE);
	}

    @Override
    public Object visit(Program program) {

        for (ICClass cls : program.getClasses()) {
            cls.accept(this);
        }

        return null;
    }

    @Override
    public Object visit(ICClass icClass) {

        for (Method m : icClass.getMethods()) {
            m.accept(this);
        }

        return null;
    }

    @Override
    public Object visit(Field field) {
        //not needed for these checks
        return null;
    }

    public Object visitMethod(Method m) {
        this.currentInitDepth = 0;
        this.currentInitContext = m;
        this.currentInitializedSymbols = new HashSet<Symbol>();

        for (Statement s : m.getStatements()) {
            s.accept(this);
        }

        Symbol retSym = m.getRetSymbolTableEntry();
        if (m.getContextType().getReturnType() != this.typeTable.getVoidType() && retSym.getInitDepth() != 0) {
            this.warnings.add(new ContextWarning("Method '" + m.getName() + "' might not return values on all code paths.", m.getLine()));
        }

        return null;
    }

    @Override
    public Object visit(VirtualMethod method) {
        return visitMethod(method);
    }

    @Override
    public Object visit(StaticMethod method) {
        return visitMethod(method);
    }

    @Override
    public Object visit(LibraryMethod method) {
        //not needed for these checks
        return null;
    }

    @Override
    public Object visit(Formal formal) {
        //not needed for these checks
        return null;
    }

    @Override
    public Object visit(PrimitiveType type) {
        //not needed for these checks
        return null;
    }

    @Override
    public Object visit(UserType type) {
        //not needed for these checks
        return null;
    }

    @Override
    public Object visit(Assignment assignment) {
        Location loc = assignment.getVariable();

        assignment.getAssignment().accept(this);

        if (loc instanceof VariableLocation) {
            VariableLocation var = (VariableLocation)loc;
            if (!var.isExternal()) {
                Symbol varSym = var.getSymbolTableEntry();
                setSymbolInit(varSym);
            }
        } else {
            //check array location for bad usage
            loc.accept(this);
        }

        return null;
    }

    @Override
    public Object visit(CallStatement callStatement) {
        return callStatement.getCall().accept(this);
    }

    @Override
    public Object visit(Return returnStatement) {
        Symbol retSym = returnStatement.getSymbolTableEntry();
        setSymbolInit(retSym);
        return null;
    }

    @Override
    public Object visit(If ifStatement) {
        ifStatement.getCondition().accept(this);

        Object prevContextBackup = this.currentInitContext;
        Set<Symbol> prevInitSymbolsBackup = this.currentInitializedSymbols;

        this.currentInitContext = ifStatement.getOperation();
        this.currentInitializedSymbols = new HashSet<Symbol>();
        this.currentInitDepth++;

        ifStatement.getOperation().accept(this);

        Set<Symbol> initializedOnAllPaths = null;

        if (ifStatement.hasElse()) {
            Set<Symbol> initOnOp = this.currentInitializedSymbols;
			
			for (Symbol sym : this.currentInitializedSymbols) {
				resetSymbolInit(sym);
			}
			
            this.currentInitContext = ifStatement.getElseOperation();
            this.currentInitializedSymbols = new HashSet<Symbol>();

            ifStatement.getElseOperation().accept(this);

            Set<Symbol> initOnElseOp = this.currentInitializedSymbols;

            //initialized on both
            initOnOp.retainAll(initOnElseOp);
            initializedOnAllPaths = initOnOp;
        }
		
		for (Symbol sym : this.currentInitializedSymbols) {
			resetSymbolInit(sym);
		}

        this.currentInitDepth--;
        this.currentInitContext = prevContextBackup;
        this.currentInitializedSymbols = prevInitSymbolsBackup;

        //set symbol init in original depth and context
        if (initializedOnAllPaths != null) {
            for (Symbol sym : initializedOnAllPaths) {
                setSymbolInit(sym);
            }
        }

        return null;
    }

    @Override
    public Object visit(While whileStatement) {
        whileStatement.getCondition().accept(this);

        Object prevContextBackup = this.currentInitContext;
        Set<Symbol> prevInitSymbolsBackup = this.currentInitializedSymbols;

        this.currentInitContext = whileStatement.getOperation();
        this.currentInitializedSymbols = new HashSet<Symbol>();
        this.currentInitDepth++;

        whileStatement.getOperation().accept(this);
		
		for (Symbol sym : this.currentInitializedSymbols) {
			resetSymbolInit(sym);
		}

        this.currentInitDepth--;
        this.currentInitContext = prevContextBackup;
        this.currentInitializedSymbols = prevInitSymbolsBackup;

        return null;
    }

    @Override
    public Object visit(Break breakStatement) {
        //not needed for these checks
        return null;
    }

    @Override
    public Object visit(Continue continueStatement) {
        //not needed for these checks
        return null;
    }

    @Override
    public Object visit(StatementsBlock statementsBlock) {
        for (Statement s : statementsBlock.getStatements()) {
            s.accept(this);
        }

        return null;
    }

    @Override
    public Object visit(LocalVariable localVariable) {
        if (localVariable.hasInitValue()) {
            localVariable.getInitValue().accept(this);

            Symbol sym = localVariable.getSymbolTableEntry();
            setSymbolInit(sym);
        }

        return null;
    }

    @Override
    public Object visit(VariableLocation location) {

        Symbol sym = location.getSymbolTableEntry();
        if (sym.getKind() == Kind.Var &&
                (this.currentInitDepth < sym.getInitDepth() ||
                (this.currentInitDepth == sym.getInitDepth() &&
                    this.currentInitContext != sym.getInitContext()))) {
            this.warnings.add(new ContextWarning("Variable '" + location.getName() + "' might not be initialized before use.", location.getLine()));
        }

        return null;
    }

    @Override
    public Object visit(ArrayLocation location) {
        location.getArray().accept(this);
        location.getIndex().accept(this);
        return null;
    }

    @Override
    public Object visit(StaticCall call) {
        for (Expression arg : call.getArguments()) {
            arg.accept(this);
        }

        return null;
    }

    @Override
    public Object visit(VirtualCall call) {
        for (Expression arg : call.getArguments()) {
            arg.accept(this);
        }

        if (call.isExternal()) {
            call.getLocation().accept(this);
        }

        return null;
    }

    @Override
    public Object visit(This thisExpression) {
        //not needed for these checks
        return null;
    }

    @Override
    public Object visit(NewClass newClass) {
        //not needed for these checks
        return null;
    }

    @Override
    public Object visit(NewArray newArray) {
        return newArray.getSize().accept(this);
    }

    @Override
    public Object visit(Length length) {
        return length.getArray().accept(this);
    }

    @Override
    public Object visit(MathBinaryOp binaryOp) {
        binaryOp.getFirstOperand().accept(this);
        binaryOp.getSecondOperand().accept(this);
        return null;
    }

    @Override
    public Object visit(LogicalBinaryOp binaryOp) {
        binaryOp.getFirstOperand().accept(this);
        binaryOp.getSecondOperand().accept(this);
        return null;
    }

    @Override
    public Object visit(MathUnaryOp unaryOp) {
        return unaryOp.getOperand().accept(this);
    }

    @Override
    public Object visit(LogicalUnaryOp unaryOp) {
        return unaryOp.getOperand().accept(this);
    }

    @Override
    public Object visit(Literal literal) {
        //not needed for these checks
        return null;
    }

    @Override
    public Object visit(ExpressionBlock expressionBlock) {
        return expressionBlock.accept(this);
    }

    public List<ContextWarning> getWarnings() {
        return warnings;
    }
}
